# import sys
# sys.path.append('/home/gpu/peng')
# from hj_dataset_devkit import SupportedDataset, load_dataset, CoordinateSystem
from __init__ import SupportedDataset, load_dataset, CoordinateSystem

def main():
    # nuscenes
    # dataset_path = '/media/gpu/sdb/datasets/nuscenes/v1.0-trainval/'
    # dataset = load_dataset(SupportedDataset.nuscenes, dataset_path)
    # hongjing
    dataset_path = '/home/gpu/peng/cache_test/'
    dataset = load_dataset(SupportedDataset.hongjing, dataset_path)
    print(f'{len(dataset.scenes)} scenes in dataset')
    # wrong_path = '/home/gpu/peng/'
    # print(dataset.check_file_structure(wrong_path))
    # dataset.load_dataset(dataset_path)
    scene = dataset.scenes[-1]
    # print(f'scene name: {scene.meta.name}')
    # print(f'scene description: {scene.meta.description}')
    # print(f'scene tags: {scene.meta.tags}')
    # print(f'scene lidar_channels: {scene.meta.lidar_channels}')
    # print(f'scene radar_channels: {scene.meta.radar_channels}')
    # print(f'scene calibration keys: {scene.calib.keys()}')
    print(f'{len(scene.frames)} frames in scene {scene.meta.name}')
    frame = scene.frames[int(len(scene.frames) / 2)]
    # print(f'lidar sensors: {frame.lidars}')
    # print(f'radar sensors: {frame.radars}')
    # print(f'camera sensors: {frame.cameras}')
    # print(f'location in world of frame: {frame.ego_pose}')
    # print(f'all sensors: {frame.all_sensors()}')
    # print(f'frame timestamp(ms): {frame.get_timestamp()}')
    print(f'{frame.cameras[-1]} timestamp(us): {frame.get_timestamp(frame.cameras[-1])}')
    coor = CoordinateSystem.world
    cloud = frame.get_merged_lidar_xyzi_cloud(coor)
    print(f'lidar cloud shape: {cloud.shape}')
    cloud = frame.get_merged_radar_cloud(coor)
    print(f'radar cloud shape: {cloud.shape}')
    img = frame.get_camera_image(frame.cameras[-1])
    print(f'camera {frame.cameras[-1]} image size: {img.width}*{img.height}')
    obstacles = frame.get_obstacles()
    print(f'{len(obstacles)} obstacles in this frame')
    ca_interval_s = 4
    frames_objs = scene.get_obstacles_with_motion(ca_interval_s)
    print(f'{len(frames_objs)} frames have tracked obstacles')
    MAP_INFO_RANGE = 100
    map_info = frame.get_map_info(coor, roi=MAP_INFO_RANGE)
    print(f'{len(map_info.lane_dividers)} lane dividers within {MAP_INFO_RANGE}m')
    print(f'{len(map_info.road_edges)} road edges within {MAP_INFO_RANGE}m')
    print(f'{len(map_info.road_markings)} road markings within {MAP_INFO_RANGE}m')
    print(f'{len(map_info.traffic_signs)} traffic signs within {MAP_INFO_RANGE}m')
    print(f'{len(map_info.traffic_lights)} traffic lights within {MAP_INFO_RANGE}m')
    # dataset.cache('../cache_test')
    # dataset.scenes[0].cache('../cache_test', force_clear=False)
    print('Test done')

if __name__ == '__main__':
    main()
